# -*- coding: utf-8 -*-
import json
import uuid
import os
import base64
import time
from typing import Optional

import aiofiles
from loguru import logger
from fastapi import Request, UploadFile, File, Form
from openai import OpenAI
from sse_starlette import EventSourceResponse

from configs import UPLOAD_IMAGE_DIR

# ==================== 配置参数 ====================
# Qwen VL API 配置
QWEN_VL_API_KEY = "sk-18c8e65919e04db690a830f94b30f73c"
QWEN_VL_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_VL_MODEL = "qwen3-vl-plus"

# 支持的图片格式
SUPPORTED_IMAGE_TYPES = {
    '.png': 'image/png',
    '.jpeg': 'image/jpeg',
    '.jpg': 'image/jpeg',
    '.webp': 'image/webp',
}


# ==================== 基础工具函数 ====================
def encode_image(image_path: str) -> str:
    """将本地图片文件编码为 Base64 字符串"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


async def call_qwen_vl_api(image_path: str, prompt_text: str, enable_thinking: bool = False, thinking_budget: int = 500):
    """
    调用 Qwen VL API 进行图片分析（流式）
    :param image_path: 本地图片路径
    :param prompt_text: 提示语
    :param enable_thinking: 是否开启思考过程
    :param thinking_budget: 推理过程最大Token数
    """
    try:
        # 检查图片格式
        ext = os.path.splitext(image_path)[1].lower()
        if ext not in SUPPORTED_IMAGE_TYPES:
            raise ValueError(f"不支持的图片格式: {ext}，只支持: {list(SUPPORTED_IMAGE_TYPES.keys())}")

        # 检查文件是否存在
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件不存在: {image_path}")

        content_type = SUPPORTED_IMAGE_TYPES[ext]
        base64_image = encode_image(image_path)
        url = f"data:{content_type};base64,{base64_image}"

        client = OpenAI(
            api_key=QWEN_VL_API_KEY,
            base_url=QWEN_VL_BASE_URL,
        )

        completion = client.chat.completions.create(
            model=QWEN_VL_MODEL,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {"url": url},
                        },
                        {"type": "text", "text": prompt_text},
                    ],
                }
            ],
            stream=True,
            extra_body={
                "enable_thinking": enable_thinking,
                "thinking_budget": thinking_budget
            } if enable_thinking else None,
        )

        for chunk in completion:
            if not chunk.choices:
                continue

            delta = chunk.choices[0].delta

            # 处理思考过程（如果启用）
            if enable_thinking and hasattr(delta, 'reasoning_content') and delta.reasoning_content:
                yield delta.reasoning_content

            # 处理正常回复内容
            if delta.content:
                yield delta.content
            else:
                yield ''

    except Exception as e:
        logger.error(f"Qwen VL API调用失败: {e}")
        import traceback
        logger.error(traceback.format_exc())
        yield ''


async def save_uploaded_image(file: UploadFile) -> str:
    """
    保存上传的图片文件到 UPLOAD_IMAGE_DIR 目录
    :param file: 上传的文件对象
    :return: 保存后的图片文件路径
    """
    # 确保目录存在
    if not os.path.exists(UPLOAD_IMAGE_DIR):
        os.makedirs(UPLOAD_IMAGE_DIR)
        logger.info(f"创建图片上传目录: {UPLOAD_IMAGE_DIR}")

    # 获取文件扩展名
    original_filename = file.filename or "image"
    if '.' in original_filename:
        file_ext = original_filename.rsplit('.', 1)[-1].lower()
        # 验证扩展名是否支持
        if f'.{file_ext}' not in SUPPORTED_IMAGE_TYPES:
            raise ValueError(f"不支持的图片格式: .{file_ext}，只支持: {list(SUPPORTED_IMAGE_TYPES.keys())}")
    else:
        file_ext = 'jpg'  # 默认扩展名

    # 生成唯一文件名
    unique_id = str(uuid.uuid4())
    timestamp = int(time.time())
    filename = f"{timestamp}_{unique_id}.{file_ext}"
    file_path = os.path.join(UPLOAD_IMAGE_DIR, filename)

    # 读取文件内容并保存
    try:
        file_content = await file.read()
        if not file_content:
            raise ValueError("上传的文件内容为空")

        async with aiofiles.open(file_path, "wb") as f:
            await f.write(file_content)

        logger.info(f"图片已保存: {file_path}")
        return file_path

    except Exception as e:
        logger.error(f"保存图片文件失败: {e}")
        raise


async def chat_iter(image_path: str, prompt_text: str, dialog_id: str, origin_query: str, token: str):
    """
    图片风险识别对话流式生成器
    :param image_path: 图片路径
    :param prompt_text: 提示语
    :param dialog_id: 对话ID
    :param origin_query: 原始查询
    :param token: 用户token
    """
    default_prompt = "帮我分析图中有什么，并分析潜在风险和隐患，给出法律依据，最后，给出整改建议"
    final_prompt = prompt_text if prompt_text else default_prompt

    try:
        res = ''
        async for content in call_qwen_vl_api(image_path, final_prompt, enable_thinking=False):
            if content is None:
                break
            res += content if content else ''
            if content:
                yield content

        final_response = json.dumps({"response": res}, ensure_ascii=False)
        yield final_response
        logger.info(f"dialog_id: {str(dialog_id)}, image_path: {image_path}, answer: {res}")

    except Exception as e:
        import traceback
        traceback.print_exc()
        error_response = json.dumps({
            "event": "error",
            "data": json.dumps({
                "code": 500,
                "data": None,
                "msg": f"图片分析失败: {str(e)}"
            }, ensure_ascii=False)
        }, ensure_ascii=False)
        yield error_response


async def risks_identification_chat(
    request: Request,
    token: str = Form(...),
    file: UploadFile = File(...),
    prompt_text: Optional[str] = Form(None),
    dialog_id: Optional[str] = Form(None),
):
    """图片风险识别助手"""
    try:
        # 保存上传的图片
        image_path = await save_uploaded_image(file)
        
        # 返回流式响应
        return EventSourceResponse(
            chat_iter(
                image_path=image_path,
                prompt_text=prompt_text or "",
                dialog_id=dialog_id or "",
                origin_query=prompt_text or "",
                token=token
            )
        )
    except Exception as e:
        import traceback
        logger.error(f"图片风险识别失败: {e}")
        logger.error(traceback.format_exc())
        # 对于同步错误，返回错误信息（SSE 需要在生成器中处理）
        error_msg = json.dumps({
            "event": "error",
            "data": json.dumps({
                "code": 500,
                "data": None,
                "msg": f"图片上传或处理失败: {str(e)}"
            }, ensure_ascii=False)
        }, ensure_ascii=False)
        
        async def error_generator():
            yield error_msg
        
        return EventSourceResponse(error_generator())

